import torch
from timm import create_model
from torchvision.datasets import ImageFolder
from torchvision import transforms
from transformers import TrainingArguments, Trainer, EvalPrediction
import numpy as np

# 1. 加载数据集
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_dataset = ImageFolder("mydata/output/melon17_full/train", transform=transform)
val_dataset = ImageFolder("mydata/output/melon17_full/val", transform=transform)

# 2. 加载EfficientNet-B0
num_classes = len(train_dataset.classes)
model = create_model("efficientnet_b0", pretrained=True, num_classes=num_classes)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(DEVICE)

# 3. 定义Trainer
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    accuracy = (predictions == labels).mean()
    return {"accuracy": accuracy}

training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    evaluation_strategy="epoch",
    num_train_epochs=10,
    logging_dir="./logs",
    learning_rate=1e-4,
    fp16=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    compute_metrics=compute_metrics,
)

# 4. 训练
trainer.train()